#pragma once

// Utility Macros & Functions
#define CHKHR_GOTO(hr, Label) \
   { if (FAILED(hr)) { wprintf(L"Error in file %S at line %d.\n", __FILE__, __LINE__); goto Label;} };

#define CHKHR_GOTO_MSG(hr, Label, wszMessage) \
	{ if (FAILED(hr)) { wprintf(L"Error in file %S at line %d.\nError Message: %s\n", __FILE__, __LINE__, wszMessage); goto Label;} };

#define CHKHR_OLEDB_GOTO(hr, Label, pItf, IID_Itf) \
	{ if (FAILED(hr)) { wprintf(L"Error in file %S at line %d.\n", __FILE__, __LINE__); DumpErrorInfo(pItf, IID_Itf); goto Label;} };

#define NUMELEM(arr) (sizeof(arr) / sizeof(arr[0])) 

// Template function that checks the NULL-ness of a COM interface and if it is non-NULL releases it & also sets it to NULL
template<class T>
void Release(T** pUnkCOMItf) {
	if (*pUnkCOMItf) {
		(*pUnkCOMItf)->Release();
		*pUnkCOMItf = NULL;
	}
}

// Utility routine for displaying OLEDB errors
void DumpErrorInfo (
    IUnknown* pObjectWithError,
    REFIID IID_InterfaceWithError
);

// COM Load/Unload Helper
class CCOMLoader {
public:
	HRESULT Load() {
		return CoInitializeEx(NULL, COINIT_MULTITHREADED);
	}
	~CCOMLoader() {
		CoUninitialize();
	}
};

// Represents an OLEDB data source, used for connection & session creation
class CSQLNCLIDataSource {
private:
	bool m_fIsConnected;
	IDBInitialize* m_pIDBInitialize;   // Data Source Initialization interface

public:
	CSQLNCLIDataSource() : m_pIDBInitialize(NULL), m_fIsConnected(false) {}

      HRESULT Connect(const wchar_t* wszDataSource, const wchar_t* wszCatalog) {
         HRESULT hr = S_OK;
         IDBProperties* pIDBProperties = NULL;

         const ULONG INIT_PROPS = 3; 
         DBPROP rgInitProps[INIT_PROPS] = {0};
         const ULONG INIT_PROPSETS = 1;
         DBPROPSET rgInitPropSets[INIT_PROPSETS] = {0};

         //Obtain access to the SQLOLEDB provider.
         hr = CoCreateInstance(CLSID_SQLNCLI10, NULL, CLSCTX_INPROC_SERVER, IID_IDBInitialize, reinterpret_cast<void **>(&m_pIDBInitialize));
         CHKHR_GOTO_MSG(hr, _Exit, L"Unable to load SQLNCLI10");

         // Set initialization property values
         SetPropertyBSTR(DBPROP_INIT_DATASOURCE, wszDataSource,  &rgInitProps[0]);
         SetPropertyBSTR(DBPROP_INIT_CATALOG,    wszCatalog, &rgInitProps[1]);
         SetPropertyBSTR(DBPROP_AUTH_INTEGRATED, L"SSPI", &rgInitProps[2]);

         // Setup the initialization property sets
         InitializePropSet(
            rgInitPropSets,
            NUMELEM(rgInitProps),
            DBPROPSET_DBINIT,
            rgInitProps);

         hr = m_pIDBInitialize->QueryInterface(IID_IDBProperties, reinterpret_cast<void**>(&pIDBProperties));
         CHKHR_GOTO_MSG(hr, _Exit, L"Failure to QI IDBInitialize.");

         hr = pIDBProperties->SetProperties(NUMELEM(rgInitPropSets), rgInitPropSets);
         CHKHR_OLEDB_GOTO(hr, _Exit, pIDBProperties, IID_IDBProperties);

         hr = m_pIDBInitialize->Initialize();
         CHKHR_OLEDB_GOTO(hr, _Exit, m_pIDBInitialize, IID_IDBInitialize);

         m_fIsConnected = true;

_Exit:
         Release(&pIDBProperties);
         CleanPropSet(rgInitPropSets);
         return hr;
      }

	HRESULT GetSession(IOpenRowset** ppIOpenRowset) {
		assert(m_pIDBInitialize);
		assert(m_fIsConnected);
		
		HRESULT hr = S_OK;
		IDBCreateSession* pIDBCreateSession = NULL;

        if (m_pIDBInitialize)
           hr = m_pIDBInitialize->QueryInterface(IID_IDBCreateSession, reinterpret_cast<void**>(&pIDBCreateSession));
		CHKHR_GOTO_MSG(hr, _Exit, L"Failure to QI IDBCreateSession.");
		
        if (pIDBCreateSession)
           hr = pIDBCreateSession->CreateSession(NULL, IID_IOpenRowset, reinterpret_cast<IUnknown**>(ppIOpenRowset));
		CHKHR_OLEDB_GOTO(hr, _Exit, pIDBCreateSession, IID_IDBCreateSession);
	_Exit:
		Release(&pIDBCreateSession);
		return hr;
	}

    ~CSQLNCLIDataSource() {
		if (m_fIsConnected) {
			assert(m_pIDBInitialize);
            HRESULT hr = S_OK;
            if (m_pIDBInitialize)
               hr = m_pIDBInitialize->Uninitialize();
			CHKHR_OLEDB_GOTO(hr, _Exit, m_pIDBInitialize, IID_IDBInitialize);
		}
	_Exit: 
		Release(&m_pIDBInitialize);
	}

private:
	void InitializePropSet( DBPROPSET* pPropSet, ULONG cProps, GUID guidPropSet, DBPROP* pProps ) {
		pPropSet->cProperties     = cProps;
		pPropSet->guidPropertySet = guidPropSet;
		pPropSet->rgProperties    = pProps;
	}

	void CleanPropSet (DBPROPSET*  pPropSet) 	{
		for (ULONG idxProp = 0; idxProp < pPropSet->cProperties; idxProp++)
			if (pPropSet->rgProperties[idxProp].vValue.vt == VT_BSTR)
				SysFreeString(pPropSet->rgProperties[idxProp].vValue.bstrVal);
	}

	void SetPropertyBSTR (DBPROPID propID, const wchar_t*  wszValue, DBPROP* pProperty) {
		pProperty->dwPropertyID    = propID;
		pProperty->dwOptions       = DBPROPOPTIONS_REQUIRED;
		pProperty->colid           = DB_NULLID;
		pProperty->vValue.vt       = VT_BSTR;
		pProperty->vValue.bstrVal  = SysAllocStringLen(wszValue, (UINT)wcslen(wszValue));
	}

	void SetPropertyBool( DBPROP* pProperty, DBPROPID dwPropID, VARIANT_BOOL boolValue) {
		pProperty->dwPropertyID     = dwPropID;
		pProperty->dwOptions        = DBPROPOPTIONS_REQUIRED;
		pProperty->colid            = DB_NULLID;
		pProperty->vValue.vt        = VT_BOOL;
		pProperty->vValue.boolVal   = boolValue;
	}

	void SetPropertyI8 ( DBPROP* pProperty, DBPROPID dwPropID, LONGLONG i8Value ) {
		pProperty->dwPropertyID     = dwPropID;
		pProperty->dwOptions        = DBPROPOPTIONS_REQUIRED;
		pProperty->colid            = DB_NULLID;
		pProperty->vValue.vt        = VT_I8;
		pProperty->vValue.llVal     = i8Value;
	}
};

class CPhotograph : public ISequentialStream {
private:
    DBREFCOUNT  m_cRef;
    BYTE*       m_pbStream;
    size_t      m_cbStreamLength;
    size_t      m_idxStreamOffset;
    
public:
    CPhotograph(size_t cbStreamLength) : m_cbStreamLength(cbStreamLength), m_idxStreamOffset(0), m_cRef(1) {
        m_pbStream = new BYTE[m_cbStreamLength];

		// Generate random data for the photograph stream
		for (size_t i = 0; i < m_cbStreamLength; i++)
			m_pbStream[i] = static_cast<BYTE>(rand() % 256);
    }
    
    ~CPhotograph() {
        delete [] m_pbStream;
    }

    STDMETHODIMP QueryInterface(REFIID riid, LPVOID* ppv) {
        if (ppv == NULL)
           return E_INVALIDARG;

        if (riid == IID_IUnknown || riid == IID_ISequentialStream)
            *ppv = reinterpret_cast<void*>(this);
        else
            *ppv = NULL;

        if (*ppv) {
            (reinterpret_cast<IUnknown*>(*ppv))->AddRef();
            return S_OK;
        }

        return E_NOINTERFACE;   
    }

    // @cmember Increments the Reference count
    STDMETHODIMP_(DBREFCOUNT) AddRef() {
        return InterlockedIncrement((long*)&m_cRef);
    }

    STDMETHODIMP_(DBREFCOUNT) Release() {
        assert(m_cRef > 0);
        ULONG cRef = InterlockedDecrement((long*) &m_cRef);

        if (!cRef)
            delete this;
        
        return cRef;
    }

    STDMETHODIMP Read(void* pBuffer, ULONG cb, ULONG* pcb) {
        if (pcb) 
           *pcb = 0;
        
        if (m_idxStreamOffset == m_cbStreamLength)
            return S_FALSE;

        size_t cbRemainingBytes = m_cbStreamLength - m_idxStreamOffset;

        if (pcb) 
           *pcb = min(cb, static_cast<ULONG>(cbRemainingBytes));
        
        memcpy(pBuffer, m_pbStream + m_idxStreamOffset,  min(cb, cbRemainingBytes)); 
        return S_OK;
    }

    STDMETHODIMP Write(const void*, ULONG, ULONG* /*pcb*/ ) {
        return E_NOTIMPL;
    }
};

void DumpErrorInfo (IUnknown*   pObjectWithError, REFIID IID_InterfaceWithError) {
    // Interfaces used in the example.
    IErrorInfo*             pIErrorInfoAll          = NULL;
    IErrorInfo*             pIErrorInfoRecord       = NULL;
    IErrorRecords*          pIErrorRecords          = NULL;
    ISupportErrorInfo*      pISupportErrorInfo      = NULL;
    ISQLErrorInfo*          pISQLErrorInfo          = NULL;
    ISQLServerErrorInfo*    pISQLServerErrorInfo    = NULL;

    // Number of error records.
    ULONG                   nRecs;
    ULONG                   nRec;

    // Basic error information from GetBasicErrorInfo.
    ERRORINFO               errorinfo;

    // IErrorInfo values.
    BSTR                    bstrDescription;
    BSTR                    bstrSource;

    // ISQLErrorInfo parameters.
    BSTR                    bstrSQLSTATE;
    LONG                    lNativeError;
 
    // ISQLServerErrorInfo parameter pointers.
    SSERRORINFO*            pSSErrorInfo    = NULL;
    OLECHAR*                pSSErrorStrings = NULL;

    // Obtain the default locale ID
    DWORD                   MYLOCALEID = GetUserDefaultLCID();

    // Only ask for error information if the interface supports it.
    if (FAILED(pObjectWithError->QueryInterface(IID_ISupportErrorInfo, (void**) &pISupportErrorInfo)))
        return;

    if (FAILED(pISupportErrorInfo->InterfaceSupportsErrorInfo(IID_InterfaceWithError)))
        return;

    // Do not test the return of GetErrorInfo. It can succeed and return
    // a NULL pointer in pIErrorInfoAll. Simply test the pointer.
    if (GetErrorInfo(0, &pIErrorInfoAll) == S_FALSE) {
        pISupportErrorInfo->Release();
        pISupportErrorInfo = NULL;
        return;
    }

    if (pIErrorInfoAll != NULL) {
        // Test to see if it's a valid OLE DB IErrorInfo interface 
        // exposing a list of records.
        if (SUCCEEDED(pIErrorInfoAll->QueryInterface(IID_IErrorRecords,(void**) &pIErrorRecords))) {
            pIErrorRecords->GetRecordCount(&nRecs);

            // Within each record, retrieve information from each
            // of the defined interfaces.
            for (nRec = nRecs - 1; (long)nRec >= 0; nRec--) {
                // From IErrorRecords, get the HRESULT and a reference
                // to the ISQLErrorInfo interface.
                pIErrorRecords->GetBasicErrorInfo(nRec, &errorinfo);
                pIErrorRecords->GetCustomErrorObject(nRec,IID_ISQLErrorInfo, (IUnknown**) &pISQLErrorInfo);

                if (pISQLErrorInfo != NULL) {
                    pISQLErrorInfo->GetSQLInfo(&bstrSQLSTATE, &lNativeError);

                    if (bstrSQLSTATE[0] == '0' && bstrSQLSTATE[1] == '1') {}
                    else {
                        // Display the SQLSTATE and native error values.
                        wprintf(L"SQLSTATE:\t%s\nNative Error:\t%ld\n",
                            bstrSQLSTATE, lNativeError);
                        
                        // SysFree BSTR references.
                        SysFreeString(bstrSQLSTATE);
                    }

                    // Get the ISQLServerErrorInfo interface from
                    // ISQLErrorInfo before releasing the reference.
                    pISQLErrorInfo->QueryInterface(IID_ISQLServerErrorInfo, (void**) &pISQLServerErrorInfo);

                    pISQLErrorInfo->Release();
                    pISQLErrorInfo = NULL;
                }

                // Test to ensure the reference is valid, then
                // get error information from ISQLServerErrorInfo.
                if (pISQLServerErrorInfo != NULL) {
                    pISQLServerErrorInfo->GetErrorInfo(&pSSErrorInfo,&pSSErrorStrings);

                    // ISQLServerErrorInfo::GetErrorInfo succeeds
                    // even when it has nothing to return. Test the
                    // pointers before using.
                    if (pSSErrorInfo) {
                        // Display the state and severity from the
                        // returned information. The error message comes
                        // from IErrorInfo::GetDescription.
                        wprintf(L"Error state:\t%d\nSeverity:\t%d\n",
                                pSSErrorInfo->bState,
                                pSSErrorInfo->bClass);

						// IMalloc::Free needed to release references
                        // on returned values. For the example, assume
                        // the g_pIMalloc pointer is valid.
                        CoTaskMemFree(pSSErrorStrings);
                        CoTaskMemFree(pSSErrorInfo);
                    }

                    pISQLServerErrorInfo->Release();
                    pISQLServerErrorInfo = NULL;
                }

                if (SUCCEEDED(pIErrorRecords->GetErrorInfo(nRec,MYLOCALEID, &pIErrorInfoRecord))) {
                    // Get the source and description (error message)
                    // from the record's IErrorInfo.
                    pIErrorInfoRecord->GetSource(&bstrSource);
                    pIErrorInfoRecord->GetDescription(&bstrDescription);

                    if (bstrSource != NULL) {
                        wprintf(L"Source:\t\t%s\n", bstrSource);
                        
                        SysFreeString(bstrSource);
                    }
                    if (bstrDescription != NULL) {
                        wprintf(L"Error message:\t%s\n", bstrDescription);
                        SysFreeString(bstrDescription);
                    }
                    pIErrorInfoRecord->Release();
                    pIErrorInfoRecord = NULL;
                }
            }

            pIErrorRecords->Release();
            pIErrorRecords = NULL;
        }
        else {
            // IErrorInfo is valid; get the source and
            // description to see what it is.
            pIErrorInfoAll->GetSource(&bstrSource);
            pIErrorInfoAll->GetDescription(&bstrDescription);

            if (bstrSource != NULL) {
                wprintf(L"Source:\t\t%s\n", bstrSource);
                SysFreeString(bstrSource);
            }
            if (bstrDescription != NULL) {
                wprintf(L"Error message:\t%s\n", bstrDescription);
                SysFreeString(bstrDescription);
            }
        }

        pIErrorInfoAll->Release();
        pIErrorInfoAll = NULL;
    }

    pISupportErrorInfo->Release();
    pISupportErrorInfo = NULL;
}